#!/usr/bin/python
#-*- coding=utf-8 -*-


from __future__ import print_function

import zmq
import json
import time
from Queue import Queue
import threading
from constant import TaskType, WorkerStatus

class Worker(threading.Thread):

    """Docstring for Worker. """


    def __init__(self, task_type, capacity, address):
        """ task_type: 处理任务类型
            capacity: 处理任务数量
            address: 服务器地址
        """

        threading.Thread.__init__(self)
        self.__tasks = Queue(capacity)
        self.__type = task_type
        self.__name = u'Worker_{}'.format(task_type.name).encode("ascii")
        self.__socket = zmq.Context().socket(zmq.DEALER)
        self.__socket.identity = self.__name
        self.__socket.connect(address)
        print(self.__name, "ready")

        rsp = {'task_type': self.__type.value, 'working': self.__tasks.qsize(), 'capacity': self.__tasks.maxsize}
        self.__socket.send_string(json.dumps(rsp))
        # print(self.__name, "send", rsp)

    def __str__(self):
        return self.__name

    def do(self, task):
        '''处理任务
        '''

        rsp = {'task_type': self.__type.value, 'working': self.__tasks.qsize(), 'capacity': self.__tasks.maxsize}
        rsp['task'] = task['task']
        rsp['idx'] = task['idx']
        rsp['working'] = self.__tasks.qsize()
        rsp['client'] = task['client']

        self.__socket.send_string(json.dumps(rsp))
        # print(self.__name, "send", rsp)

        print(self.__name, "working", task['task'])

        # TODO: 将工作逻辑写在这替代 sleep
        time.sleep(2)

        print(self.__name, "finish", rsp['task'])


    def check_task(self):
        """: 检查是否有任务完成
        """
        while True:
            for i in xrange(self.__tasks.qsize()):
                thread = self.__tasks.get()
                if (thread.isAlive()):
                    self.__tasks.put(thread)
                else:
                    # print(thread, 'thread finish')
                    rsp = {'task_type': self.__type.value, 'working': self.__tasks.qsize(), 'capacity': self.__tasks.maxsize}
                    self.__socket.send_string(json.dumps(rsp))

            #1s检查一次
            time.sleep(1)


    def run(self):

        thread = threading.Thread(target = (lambda worker: worker.check_task()), args = (self, ))
        thread.setDaemon(True)
        thread.start()

        while True:
            req = self.__socket.recv()
            req = json.loads(req)
            # print("{}: recv  {}".format(self.__name, req), self.__tasks.qsize())

            if (self.__tasks.full()):
                # rsp = {'task_type': self.__type.value, 'working': self.__tasks.qsize(), 'capacity': self.__tasks.maxsize}
                # self.__socket.send_string(json.dumps(rsp))
                # print(self.__name, rsp)
                continue

            thread = threading.Thread(target = (lambda worker, task: worker.do(task)), args = (self, req, ))
            self.__tasks.put(thread)

            thread.setDaemon(True)
            thread.start()



#例子
def main():
    """TODO: Docstring for main.
    :returns: TODO

    """
    worker_a = Worker(TaskType.A, 1, 'tcp://localhost:50001')
    worker_b = Worker(TaskType.B, 1, 'tcp://localhost:50001')
    worker_a.start()
    worker_b.start()
    worker_a.join()
    worker_b.join()


if __name__ == "__main__":
    main()
